package com.lvhx.springboot.partitioner;


import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.core.BatchStatus;
import org.springframework.batch.core.ExitStatus;
import org.springframework.batch.core.Step;
import org.springframework.batch.core.StepExecution;
import org.springframework.batch.core.partition.StepExecutionSplitter;
import org.springframework.batch.core.partition.support.TaskExecutorPartitionHandler;
import org.springframework.core.task.SyncTaskExecutor;
import org.springframework.core.task.TaskExecutor;
import org.springframework.core.task.TaskRejectedException;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;

import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;

public class CustomTaskExecutorPartitionHandler extends TaskExecutorPartitionHandler {
    private static Logger logger = LoggerFactory.getLogger(CustomTaskExecutorPartitionHandler.class);

    private Step step;

    private TaskExecutor taskExecutor = new SyncTaskExecutor();

    public CustomTaskExecutorPartitionHandler() {
    }


    public Step getStep() {
        return step;
    }

    public void setStep(Step step) {
        super.setStep(step);
        this.step = step;
    }

    public TaskExecutor getTaskExecutor() {
        return taskExecutor;
    }

    public void setTaskExecutor(TaskExecutor taskExecutor) {
        super.setTaskExecutor(taskExecutor);
        this.taskExecutor = taskExecutor;
    }

    @Override
    protected Set<StepExecution> doHandle(StepExecution masterStepExecution, Set<StepExecution> partitionStepExecutions)
            throws Exception {
        Assert.notNull(step, "A Step must be provided.");
        final Set<Future<StepExecution>> tasks = new HashSet<Future<StepExecution>>(getGridSize());
        final Set<StepExecution> result = new HashSet<StepExecution>();

        try {
        } catch (Throwable th) {
            logger.warn("get EventContext in current thread exception.", th);
        }

        for (final StepExecution stepExecution : partitionStepExecutions) {

            final FutureTask<StepExecution> task = createTask(step, stepExecution);

            try {
                taskExecutor.execute(task);
                tasks.add(task);
            } catch (TaskRejectedException e) {
                // couldn't execute one of the tasks
                ExitStatus exitStatus = ExitStatus.FAILED
                        .addExitDescription("TaskExecutor rejected the task for this step.");
                /*
                 * Set the status in case the caller is tracking it through the
                 * JobExecution.
                 */
                stepExecution.setStatus(BatchStatus.FAILED);
                stepExecution.setExitStatus(exitStatus);
                result.add(stepExecution);
            }
        }

        for (Future<StepExecution> task : tasks) {
            result.add(task.get());
        }

        return result;
    }

    @Override
    public Collection<StepExecution> handle(StepExecutionSplitter stepExecutionSplitter,
                                            StepExecution masterStepExecution) throws Exception {

        masterStepExecution.getExecutionContext().put("HANDLER_TYPE", CustomTaskExecutorPartitionHandler.class.getName());
        return super.handle(stepExecutionSplitter, masterStepExecution);
    }

    /**
     * Creates the task executing the given step in the context of the given
     * execution.
     *
     * @param step          the step to execute
     * @param stepExecution the given execution
     * @return the task executing the given step
     */
    protected FutureTask<StepExecution> createTask(final Step step, final StepExecution stepExecution) {
        return new FutureTask<StepExecution>(new Callable<StepExecution>() {
            @Override
            public StepExecution call() throws Exception {

                String jobName = null;
                String sharding = null;
                try {
                    jobName = stepExecution.getJobExecution().getJobInstance().getJobName();
                    if (StringUtils.isEmpty(jobName))
                        jobName = stepExecution.getJobParameters().getString("JOB_NAME");
                } catch (Throwable th) {
                    logger.error("get Step Context exception.", th);
                }


                step.execute(stepExecution);

                return stepExecution;
            }
        });
    }

}
